#include "timer.h"

#include <linux/fs.h>
#include <linux/cdev.h>
#include <linux/device.h>
#include <linux/sched.h>
#include <linux/timer.h>
#include <linux/wait.h> // For wait queues
#include <linux/uaccess.h>

#include "log.h"

#define DEVICE_NAME "jit_timer"
#define CLASS_NAME  "jit"
#define BYTES_PER_LINE 128

int timer_interval = 1000;

static dev_t jit_dev_t;
static struct cdev jit_cdev;
static struct class *jit_class;

struct jit_timer_context {
    struct timer_list timer;
    wait_queue_head_t wq;
    atomic_t counter; // Use atomic counter for safe cross-CPU access
};

// Timer context is now global, as it's tied to the device, not a single read() call
static struct jit_timer_context g_jit_ctx;

// Timer callback: MINIMAL WORK. No snprintf, no logging.
static void jit_timer_work(struct timer_list *timer) {
    atomic_inc(&g_jit_ctx.counter); // Safely increment counter
    wake_up_interruptible(&g_jit_ctx.wq); // Wake up sleeping process
}

static ssize_t jit_timer_read(struct file *filp, char __user *ubuf, size_t count, loff_t *f_pos) {
    int loops = 5;
    int loop_count = 0;
    char *buf, *p_buf;
    size_t buf_size, total_len = 0;
    ssize_t ret = 0;
    ulong prev_jiffies;
    int cpu;

    if (*f_pos > 0) {
        return 0;
    }

    buf_size = (BYTES_PER_LINE * (loops + 2));
    buf = kzalloc(buf_size, GFP_KERNEL);
    if (!buf) return -ENOMEM;
    p_buf = buf;

    // Prepare header in process context
    p_buf += sprintf(p_buf, "time\t\tdelta\tinirq\tpid\tcpu\tcommand\n");
    prev_jiffies = jiffies;

    // Setup timer
    timer_setup(&g_jit_ctx.timer, jit_timer_work, 0);
    g_jit_ctx.timer.expires = jiffies + msecs_to_jiffies(timer_interval);
    add_timer(&g_jit_ctx.timer);

    for (loop_count = 0; loop_count < loops; loop_count++) {
        // Go to sleep and wait for the timer to wake us up.
        // wait_event_interruptible is the modern, safe way to do this.
        wait_event_interruptible(g_jit_ctx.wq, atomic_read(&g_jit_ctx.counter) > loop_count);

        // Woken up! Now we do the data formatting in process context.
        ulong now = jiffies;
        cpu = get_cpu();
        p_buf += sprintf(p_buf, "%lu\t%lu\t%d\t%d\t%d\t%s\n",
            now, now - prev_jiffies, in_interrupt() ? 1 : 0, task_pid_nr(current), cpu, current->comm);
        prev_jiffies = now;
        put_cpu();

        // Reschedule the timer for the next loop
        if (loop_count < loops - 1) {
            g_jit_ctx.timer.expires += msecs_to_jiffies(timer_interval);
            mod_timer(&g_jit_ctx.timer, g_jit_ctx.timer.expires);
        }
    }

    total_len = p_buf - buf;
    if (copy_to_user(ubuf, buf, total_len)) {
        ret = -EFAULT;
    } else {
        ret = total_len;
        *f_pos += total_len;
    }

    kfree(buf);
    del_timer_sync(&g_jit_ctx.timer); // Clean up timer
    return ret;
}

const struct file_operations jit_timer_fops = {
    .owner = THIS_MODULE,
    .read = jit_timer_read,
    .llseek = no_llseek,
};

int jit_timer_init(void) {
    int ret;

    // Initialize wait queue
    init_waitqueue_head(&g_jit_ctx.wq);
    atomic_set(&g_jit_ctx.counter, 0);

    ret = alloc_chrdev_region(&jit_dev_t, 0, 1, DEVICE_NAME);
    if (ret < 0) return ret;

    jit_class = class_create(CLASS_NAME);
    if (IS_ERR(jit_class)) {
        unregister_chrdev_region(jit_dev_t, 1);
        return PTR_ERR(jit_class);
    }

    device_create(jit_class, NULL, jit_dev_t, NULL, DEVICE_NAME);

    cdev_init(&jit_cdev, &jit_timer_fops);
    ret = cdev_add(&jit_cdev, jit_dev_t, 1);
    if (ret < 0) {
        device_destroy(jit_class, jit_dev_t);
        class_destroy(jit_class);
        unregister_chrdev_region(jit_dev_t, 1);
        return ret;
    }

    LOGI("timer module loaded, device created at /dev/%s", DEVICE_NAME);
    return 0;
}

void jit_timer_exit(void) {
    del_timer_sync(&g_jit_ctx.timer);
    cdev_del(&jit_cdev);
    device_destroy(jit_class, jit_dev_t);
    class_destroy(jit_class);
    unregister_chrdev_region(jit_dev_t, 1);
    LOGI("timer module unloaded");
}